use std::path::Path;

use cxx::UniquePtr;
use serde::{Deserialize, Serialize};

#[cxx::bridge]
pub mod ffi {

    #[namespace = "chatglm"]
    unsafe extern "C++" {
        //include!("chatglm.h");
        include!("extern.h");
        pub type Pipeline;
        pub type ChatMessage;
        pub type GenerationConfig;
        pub type BaseStreamer;
        pub type ExternStreamer;
        pub type ExternPipeline;
        pub fn extern_pipeline_new(
            path: &str,
            gen_config: UniquePtr<GenerationConfig>,
        ) -> UniquePtr<ExternPipeline>;
        pub fn generate_string(self: &ExternPipeline, prompt: &str, use_callback: bool) -> String;
        pub fn set_callback(self: &ExternPipeline, process_callback: fn(&str), end_callback: fn());
        pub fn generation_config_new(
            max_length: i32,
            max_new_tokens: i32,
            max_context_length: i32,
            do_sample: bool,
            top_k: i32,
            top_p: f32,
            temperature: f32,
            repetition_penalty: f32,
            num_threads: i32,
        ) -> UniquePtr<GenerationConfig>;
    }
}

#[derive(Debug,Clone,Serialize,Deserialize)]
pub struct GenerationConfig {
    pub max_length: i32,
    pub max_new_tokens: i32,
    pub max_context_length: i32,
    pub do_sample: bool,
    pub top_k: i32,
    pub top_p: f32,
    pub temperature: f32,
    pub repetition_penalty: f32,
    pub num_threads: i32,
}

impl Default for GenerationConfig {
    fn default() -> Self {
        Self {
            max_length: 2048,
            max_new_tokens: -1,
            max_context_length: 512,
            do_sample: true,
            top_k: 0,
            top_p: 0.7,
            temperature: 0.95,
            repetition_penalty: 1f32,
            num_threads: 0,
        }
    }
}

impl GenerationConfig {
    pub fn into_ffi_type(&self) -> UniquePtr<ffi::GenerationConfig> {
        ffi::generation_config_new(
            self.max_length,
            self.max_new_tokens,
            self.max_context_length,
            self.do_sample,
            self.top_k,
            self.top_p,
            self.temperature,
            self.repetition_penalty,
            self.num_threads,
        )
    }
}
pub struct Pipeline {
    inner: UniquePtr<ffi::ExternPipeline>,
}

impl Pipeline {
    pub fn new<P>(path : P ,config : &GenerationConfig) -> anyhow::Result<Self> where P : AsRef<Path> {
        let ffi_generate_config = config.into_ffi_type();
        let path = path.as_ref().to_str().ok_or(anyhow::Error::msg("Failed to convert path to &str"))?;
        Ok(Self {
            inner : ffi::extern_pipeline_new(path, ffi_generate_config)
        })
    }
    pub fn generate(&self,prompt : &str) -> String {
        self.inner.generate_string(prompt, false)
    }
    pub fn generate_with_callback(&self,prompt: &str, process_callback: fn(&str), end_callback: fn()) -> String {
        self.inner.set_callback(process_callback, end_callback);
        self.inner.generate_string(prompt, true)
    }
}


#[cfg(test)]
mod tests {
    use super::*;
    const MODEL_PATH : &str = "model/chatglm-ggml-q4_1.bin";
    #[test]
    fn it_works() -> anyhow::Result<()> {
        let pipeline = Pipeline::new(MODEL_PATH, &GenerationConfig::default())?;
        let res = pipeline.generate("hello");
        println!("{}",res);
        Ok(())
    }
}
